import os
import cv2
from speed_test import getSpeed
from occlusion import DetectOcclusion

base_path = r'C:\Users\29626\Desktop\2022-02-11'
de_oc = DetectOcclusion()

def forfile(path):
    for file in os.listdir(path):
        abs_file = path + '/' + file
        if file.endswith(".mp4"):
            # print(file)
            try:
                if file.split("_")[2] in ['01', '09']:
                    detect(abs_file)
            except Exception as e:
                pass
        elif os.path.isdir(abs_file):
            forfile(abs_file)

def detect(video):
    global de_oc
    vc = cv2.VideoCapture(video)  # 读入视频文件
    de_oc.set(stride=1,  # stride: 输入步长(秒)，每1s输入一帧做差
              during=90)  # during: 总持续时间(秒)，每个视频遮挡检测持续300s
    de_oc.setSaveName(video)

    for i in range(9999):
        (ret, frame) = vc.read()
        if not ret:  # 如果没有获取到数据，则结束循环
            break
        if i % (de_oc.stride * 25) == 0:  # 间隔帧 = 输入步长×视频帧率
            speed = getSpeed(frame)
            # print("列车当前速度: ", speed, "(km/h)")
            de_oc.next(frame, speed)  # if speed == 0 -------- return


if __name__ == '__main__':
    forfile(base_path)
